#include <ros/ros.h> 
#include <image_transport/image_transport.h> 
#include <cv_bridge/cv_bridge.h> 
#include <sensor_msgs/image_encodings.h> 
#include <opencv2/core.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>
using namespace cv;
using namespace std;

// 全局HSV阈值变量
int yellow_low_h = 15, yellow_high_h = 30;
int yellow_low_s = 100, yellow_high_s = 255;
int yellow_low_v = 100, yellow_high_v = 255;

// 黑色道线处理参数
int black_threshold = 40;          // 自适应阈值参数
int canny_threshold = 30;          // 边缘检测阈值
int morph_size = 5;                // 形态学核大小
int roi_height = 300;              // 关注区域高度

// 中心线拟合参数
const int MAX_POINTS = 20; // 最大历史中心点数
int poly_degree = 2;      // 多项式阶数

// 创建控制窗口和滑动条
void createControlWindow() {
    cv::namedWindow("Control", WINDOW_AUTOSIZE);
    
    // 黄色车道线阈值滑动条
    cv::createTrackbar("YLowH", "Control", &yellow_low_h, 179);
    cv::createTrackbar("YHighH", "Control", &yellow_high_h, 179);
    cv::createTrackbar("YLowS", "Control", &yellow_low_s, 255);
    cv::createTrackbar("YHighS", "Control", &yellow_high_s, 255);
    cv::createTrackbar("YLowV", "Control", &yellow_low_v, 255);
    cv::createTrackbar("YHighV", "Control", &yellow_high_v, 255);
    
    // 黑色道线参数
    cv::createTrackbar("BlackThresh", "Control", &black_threshold, 100);
    cv::createTrackbar("CannyThresh", "Control", &canny_threshold, 100);
    cv::createTrackbar("MorphSize", "Control", &morph_size, 20);
    cv::createTrackbar("ROIHeight", "Control", &roi_height, 600);
    
    // 中心线参数
    cv::createTrackbar("Poly Degree", "Control", &poly_degree, 3);
}

void imageCallback(const sensor_msgs::ImageConstPtr& msg) {
    // 1. 转换ROS图像为OpenCV格式
    cv_bridge::CvImagePtr cv_ptr;
    try {
        cv_ptr = cv_bridge::toCvCopy(msg, "bgr8");
    } catch (...) {
        ROS_ERROR("Image conversion failed");
        return;
    }
    Mat frame = cv_ptr->image;
    
    // 设置关注区域（只处理图像下方）
    Rect roi_rect(0, frame.rows - roi_height, frame.cols, roi_height);
    Mat roi_frame = frame(roi_rect);

    // 2. 使用当前滑动条值生成阈值
    Scalar yellow_lower(yellow_low_h, yellow_low_s, yellow_low_v);
    Scalar yellow_upper(yellow_high_h, yellow_high_s, yellow_high_v);

    // 3. 预处理
    Mat hsv, mask_yellow;
    GaussianBlur(roi_frame, roi_frame, Size(5, 5), 0);
    cvtColor(roi_frame, hsv, COLOR_BGR2HSV);

    // 4. 黄色道线分割（保持不变）
    inRange(hsv, yellow_lower, yellow_upper, mask_yellow);
    Mat kernel = getStructuringElement(MORPH_RECT, Size(15, 15));
    morphologyEx(mask_yellow, mask_yellow, MORPH_CLOSE, kernel);

    // 5. 改进的黑色道线处理
    Mat gray, mask_black, edges, lab;
    cvtColor(roi_frame, gray, COLOR_BGR2GRAY);
    
    // 自适应阈值处理
    adaptiveThreshold(gray, mask_black, 255, ADAPTIVE_THRESH_GAUSSIAN_C, 
                     THRESH_BINARY_INV, 2*black_threshold+3, 5);
    
    // 边缘检测增强
    Canny(gray, edges, canny_threshold, canny_threshold*2);
    mask_black = mask_black | edges;
    
    // 在LAB空间处理黑色道线（提高亮度不变性）
    cvtColor(roi_frame, lab, COLOR_BGR2Lab);
    vector<Mat> lab_channels;
    split(lab, lab_channels);
    Mat L = lab_channels[0];
    
    // 对L通道进行自适应阈值
    Mat lab_mask;
    adaptiveThreshold(L, lab_mask, 255, ADAPTIVE_THRESH_GAUSSIAN_C, 
                     THRESH_BINARY_INV, 2*black_threshold+5, 7);
    mask_black = mask_black & lab_mask;
    
    // 动态形态学操作
    int morph_value = max(3, morph_size);
    Mat morph_kernel = getStructuringElement(MORPH_RECT, Size(morph_value, morph_value));
    morphologyEx(mask_black, mask_black, MORPH_CLOSE, morph_kernel);
    morphologyEx(mask_black, mask_black, MORPH_OPEN, morph_kernel);

    // 6. 车道线定位
    Mat result = roi_frame.clone();
    vector<vector<Point>> contours;
    
    // 存储左右车道线矩形
    vector<Rect> left_rects, right_rects;
    int img_center_x = roi_frame.cols / 2;
    
    // 黄色车道线（左侧）
    findContours(mask_yellow, contours, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE);
    for (const auto& cnt : contours) {
        if (contourArea(cnt) > 500) {
            Rect rect = boundingRect(cnt);
            // 根据水平位置分类
            if (rect.x + rect.width/2 < img_center_x) { 
                left_rects.push_back(rect);
                rectangle(result, rect, Scalar(0, 255, 255), 2);
            }
        }
    }
    
    // 黑色车道线（右侧）
    findContours(mask_black, contours, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE);
    for (const auto& cnt : contours) {
        if (contourArea(cnt) > 500) {
            Rect rect = boundingRect(cnt);
            if (rect.x + rect.width/2 > img_center_x) {
                right_rects.push_back(rect);
                rectangle(result, rect, Scalar(255, 0, 0), 2);
            }
        }
    }
    
    // 7. 中心线计算
    static vector<Point> center_points; // 存储历史中心点
    Point curr_center(-1, -1);
    
    if (!left_rects.empty() && !right_rects.empty()) {
        // 选择最底部的左右车道线（最近点）
        auto left_rect = *max_element(left_rects.begin(), left_rects.end(),
            [](const Rect& a, const Rect& b) { return a.y < b.y; });
        
        auto right_rect = *max_element(right_rects.begin(), right_rects.end(),
            [](const Rect& a, const Rect& b) { return a.y < b.y; });
        
        // 计算左右基准点（矩形底部中心）
        Point left_base(
            left_rect.x + left_rect.width/2,
            left_rect.y + left_rect.height
        );
        
        Point right_base(
            right_rect.x + right_rect.width/2,
            right_rect.y + right_rect.height
        );
        
        // 计算当前中心点
        curr_center = Point(
            (left_base.x + right_base.x)/2,
            (left_base.y + right_base.y)/2
        );
        
        // 保存历史中心点
        center_points.push_back(curr_center);
        if (center_points.size() > MAX_POINTS) {
            center_points.erase(center_points.begin());
        }
        
        // 多项式拟合中心线（使用OpenCV内置的最小二乘法）
        if (center_points.size() >= 3) {
            // 确保阶数有效
            if (poly_degree < 1) poly_degree = 1;
            else if (poly_degree > 3) poly_degree = 3;
            
            int numPoints = center_points.size();
            int numCoefficients = poly_degree + 1;
            
            // 构建矩阵A（自变量y的幂次）和向量B（因变量x）
            Mat A(numPoints, numCoefficients, CV_64F);
            Mat B(numPoints, 1, CV_64F);
            
            for (int i = 0; i < numPoints; i++) {
                double y = center_points[i].y;
                double x = center_points[i].x;
                
                for (int j = 0; j < numCoefficients; j++) {
                    A.at<double>(i, j) = pow(y, j); // y^0, y^1, y^2...
                }
                B.at<double>(i, 0) = x; // x值
            }
            
            // 求解最小二乘问题 (使用SVD分解确保数值稳定性)
            Mat coeffs;
            solve(A, B, coeffs, DECOMP_SVD);
            
            // 生成拟合曲线上的点
            vector<Point> fitted_points;
            for (int y = roi_frame.rows-1; y >= roi_frame.rows/2; y -= 5) {
                double x = 0;
                for (int j = 0; j < numCoefficients; j++) {
                    x += coeffs.at<double>(j, 0) * pow(y, j);
                }
                if (x >= 0 && x < roi_frame.cols) {
                    fitted_points.push_back(Point(x, y));
                }
            }
            
            // 绘制拟合中心线（红色）
            if (fitted_points.size() > 1) {
                polylines(result, fitted_points, false, Scalar(0, 0, 255), 3);
            }
        }
        
        // 绘制当前中心点（绿色实心圆）
        circle(result, curr_center, 8, Scalar(0, 255, 0), -1);
        
        // 计算并显示偏移量
        int offset = curr_center.x - img_center_x;
        putText(result, "Offset: " + to_string(offset) + "px", 
                Point(20, 40), FONT_HERSHEY_SIMPLEX, 0.8, Scalar(0,200,200), 2);
    }

    // 8. 创建掩膜可视化
    Mat mask_display;
    vector<Mat> channels = {mask_yellow, mask_black, Mat::zeros(roi_frame.size(), CV_8U)};
    merge(channels, mask_display);

    // 9. 显示处理结果
    // 在全尺寸图像上绘制ROI区域
    Mat full_display = frame.clone();
    rectangle(full_display, roi_rect, Scalar(0, 255, 0), 2);
    
    imshow("Original", full_display);
    imshow("Mask", mask_display);
    imshow("Result", result);
    waitKey(10);
}

int main(int argc, char** argv) {
    ros::init(argc, argv, "lane_detection_node");
    ros::NodeHandle nh;
    
    // 创建OpenCV控制窗口
    createControlWindow();
    namedWindow("Original");
    namedWindow("Mask");
    namedWindow("Result");

    // 订阅图像话题
    image_transport::ImageTransport it(nh);
    image_transport::Subscriber sub = it.subscribe("/camera/color/image_raw", 1, imageCallback);

    ros::Rate loop_rate(30);
    while (ros::ok()) {
        ros::spinOnce();
        loop_rate.sleep();
    }
    
    // 关闭所有OpenCV窗口
    destroyAllWindows();
    return 0;
}